package com.weilus.gateway;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.context.ApplicationContext;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.PathContainer;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.util.pattern.PathPattern;
import org.springframework.web.util.pattern.PathPatternParser;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * spring:
 *   cloud:
 *     gateway:
 *       routes:
 *         - id: flux-call
 *           filters:
 *           - name: TokenRateLimiter
 *             args:
 *             - GET /test/**=5/s           #5次每秒
 *             - /test/other=100/s          #100次每秒
 */
@Component
@Slf4j
public class TokenRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<TokenRateLimiterGatewayFilterFactory.Config> {

    private ApplicationContext context;
    public TokenRateLimiterGatewayFilterFactory(ApplicationContext context) {
        super(TokenRateLimiterGatewayFilterFactory.Config.class);
        this.context = context;
    }

    @Override
    public GatewayFilter apply(TokenRateLimiterGatewayFilterFactory.Config config) {
        return (exchange, chain) -> {
            PathContainer path = PathContainer.parsePath(exchange.getRequest().getURI().getRawPath());
            HttpMethod method = exchange.getRequest().getMethod();
            String token = Optional.ofNullable(exchange.getRequest().getHeaders().getFirst("token"))
                    .orElse(exchange.getRequest().getQueryParams().getFirst("token"));
            if(!StringUtils.isEmpty(token)) {
                Set<LimiterMatcher> matches = config.getAuthMatchers()
                        .stream()
                        .filter(m -> m.getPathPattern().matches(path) && (m.getMethod() == null || m.getMethod() == method))
                        .collect(Collectors.toSet());
                if (!CollectionUtils.isEmpty(matches)) {
                    LimiterMatcher m = matches.toArray(new LimiterMatcher[]{})[0];
                    RedisRateLimiter limiter = new RedisRateLimiter(m.getLimit(), m.getLimit());
                    limiter.setApplicationContext(context);
                    return limiter.isAllowed(null, token + path).flatMap((response) -> {
                        response.getHeaders().forEach((h, v) -> exchange.getResponse().getHeaders().add(h, v));
                        if (response.isAllowed()) {
                            return chain.filter(exchange);
                        } else {
                            log.debug("{} {} 限流 {}/s",m.getMethod(),m.getPathPattern(),m.getLimit());
                            ServerWebExchangeUtils.setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
                            return exchange.getResponse().setComplete();
                        }
                    });
                }
            }
            return chain.filter(exchange);
        };
    }

    public static class Config extends ArrayList<String> {
        private Set<LimiterMatcher> authMatchers = new HashSet<>();

        @Override
        public boolean add(String o) {
            authMatchers.add(new LimiterMatcher(o));
            return super.add(o);
        }

        public Set<LimiterMatcher> getAuthMatchers() {
            return authMatchers;
        }
    }

    @Data
    public static class LimiterMatcher {
        public final static Pattern timesPattern = Pattern.compile("^(\\d{1,}?)/[s,m,h,d]$");
        public final static PathPatternParser pathPatternParser = new PathPatternParser();

        private PathPattern pathPattern;
        private HttpMethod method;
        private int limit;

        public LimiterMatcher(String config) {
            String[] arr = StringUtils.delimitedListToStringArray(config, "=");
            String pathConfig = arr[0], authConfig = arr[1];
            String[] pathArr = StringUtils.delimitedListToStringArray(pathConfig, " ");
            this.pathPattern = pathPatternParser.parse(pathArr.length == 1 ? pathArr[0] : pathArr[1]);
            this.method = HttpMethod.resolve(pathArr.length == 1 ? null : pathArr[0]);
            Matcher timesMatcher = timesPattern.matcher(authConfig);
            if (timesMatcher.find()) {
                this.limit = Integer.valueOf(timesMatcher.group(1));
            }
        }

        @Override
        public String toString() {
            return "{" +
                    "method=" + method + ",pathPattern=" + pathPattern + ", limit=" + limit +
                    '}';
        }
    }
}